# -*- coding: utf-8 -*-
"""
Create Time:  2020/4/25 10:19
Author:       jinas
Email:        jinasuo@163.com
Description:  使用tensorflow 实现MNIST手写数据识别网络
本文使用的模型：LeNet
2层卷积，2层最大池化层，2层全连层
卷积核，全网使用5*5 ，padding="VALID" ,激活函数是Relu
"""

import numpy as np
import tensorflow as tf
from matplotlib import pyplot as plt
from tensorflow_core.examples.tutorials.mnist import input_data

tf.logging.set_verbosity(tf.logging.INFO)

#加载数据
mnist = input_data.read_data_sets("./dataset/MNIST/",one_hot=True)
'''
mnists数据组成：
1、train
（1）images 55000*784（28*28）
（2）labels  55000*10
2、validation 
（1）images  5000*784
（2）labels  5000*10
3、test
（1）images  10000*784
（2）labels  10000*10
'''
#先看看图片样式
plt.figure(figsize=(8,8))
for idx in range(16):
    plt.subplot(4,4,idx+1)
    plt.axis('off')
    plt.title("[{}]".format(np.argmax(mnist.train.labels[idx])))
    plt.imshow(mnist.train.images[idx].reshape(28,28))
plt.show()

#开始搭建网路模型
#设置输入和输出的占位符
x = tf.placeholder("float",[None,784],name="x")
y = tf.placeholder("float",[None,10],name="y")

#将输入转为4维度的，分别是，图像数量，高宽，通道数
x_image = tf.reshape(x,[-1,28,28,1])

#第一个卷积层
with tf.name_scope("conv1"):
    #该方法输入分别是：被卷积对象，卷积核数量（即输出的通道数），卷积核大小，填充方式，激活函数
    C1 = tf.contrib.slim.conv2d(x_image,6,[5,5],padding="VALID",activation_fn=tf.nn.relu)
    #输入是 28*28*1，卷积核是5*5，不填充,stride=1，卷积核数量是6，结果输出是24*24*6，{(28-5+0)/1 + 1 = 24}
#第二层：最大池化层，池化不对通道数进行改变
with tf.name_scope("pool1"):
    #该方法输入分别是：被池化对象，池化核大小，步长，填充方式
    S2 = tf.contrib.slim.max_pool2d(C1,[2,2],stride=[2,2],padding="VALID")
    #输入是24*24*6，输出是12*12*6{（24-2+0）/2 + 1}
#第三层：卷积层
with tf.name_scope("conv2"):
    C3 = tf.contrib.slim.conv2d(S2,16,[5,5],padding="VALID",activation_fn=tf.nn.relu)
    #输入是12*12*6，输出是8*8*16  （12-5+0）/1 + 1 =8
#第四层：最大池化层
with tf.name_scope("pool2"):
    S4 = tf.contrib.slim.max_pool2d(C3,[2,2],stride=[2,2],padding="VALID")
    #输入是8*8*16 ，输出是4*4*16     （8-2+0）/2 + 1= 4
#第五层：全连层
with tf.name_scope("fc1"):
    #flatten----保留第一维度，将其余剩下维度展开成一个行向量，返回值是一个二维数据，[第一维度，子张量乘积]
    #e.g. S4 -- X*4*4*16 ，--> flatten --> X*(4*4*16)
    S4_flat = tf.contrib.slim.flatten(S4)
    #该方法对应操作是全连接层，输入分别是：操作对象，维度，激活函数
    C5 = tf.contrib.slim.fully_connected(S4_flat,120,activation_fn=tf.nn.relu)
    #输入是X*256，输出是X*120
#第六层：全连层
with tf.name_scope("fc2"):
    F6 = tf.contrib.slim.fully_connected(C5,84,activation_fn=tf.nn.relu)
    #输入是X*120，输出是X*84
#添加dropout操作，随机丢弃部分数据，提高推光能力，减少过拟合
with tf.name_scope("dropout"):
    keep_prob = tf.placeholder(tf.float32,name="keep_prob")
    #该方法输入是：操作对象，保留概率。
    #理解：以一定概率保留数据，保留数据变为x/keep_prob，如果不保留就变为0. 有keep_prob比例的数据保留下来
    F6_drop = tf.nn.dropout(F6,keep_prob)
with tf.name_scope("fc3"):
    logits = tf.contrib.slim.fully_connected(F6_drop,10,activation_fn=None)
    #输入是：X*84，输出是X*10，没有激活函数
#定义模型的损失函数，优化器等
#softmax_cross_entropy_with_logits--输入：logits--是概率替代物，--> softmax --> 概率，然后与labels进行交叉熵计算；
#labels -- 标签，是概率，类似独热编码，但是这是指每个类别（索引）下的概率，每一行是概率分布，概率值和为0
#如果labels的每一行是one-hot表示，也就是只有一个地方为1（或者说100%），其他地方为0（或者说0%），还可以使用tf.sparse_softmax_cross_entropy_with_logits()。之所以用100%和0%描述，就是让它看起来像一个概率分布。
'''
reduce_mean -- 求平均值
参数axis--不指定（计算所有元素的平均值），0--计算每列的平均值，1--计算每行的平均值
'''
cross_entropy_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits,labels=y))
'''
add_n -- 将列表中的每个元素，对应位置相加
返回 -- 列表中元素一样size的数据
'''
l2_loss = tf.add_n([
    tf.nn.l2_loss(w) for w in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
])

for w in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES):
    print(w.name)
    '''
    tf.summary.histogram -- 显示直方图信息
    输入分别是： tags--标签；values--值
    扩展：
    （1）summary--所有需要在TensorBoard上展示的统计结果。
    （2）tf.name_scope()：为Graph中的Tensor添加层级，TensorBoard会按照代码指定的层级进行展示，初始状态下只绘制最高层级的效果，点击后可展开层级看到下一层的细节。
    （3）tf.summary.scalar()：添加标量统计结果。
    （4）tf.summary.histogram()：添加任意shape的Tensor，统计这个Tensor的取值分布。
    （5）tf.summary.merge_all()：添加一个操作，代表执行所有summary操作，这样可以避免人工执行每一个summary op。
    （6）tf.summary.FileWrite：用于将Summary写入磁盘，需要制定存储路径logdir，如果传递了Graph对象，则在Graph Visualization会显示Tensor Shape Information。执行summary op后，将返回结果传递给add_summary()方法即可。
        可以调用其add_summary()方法将训练过程数据保存在filewriter指定的文件中。
    '''
    tf.summary.histogram(w.name,w)
'''
目标函数 = 损失函数 + L2正则
这里超参数设置为 7e-5 ，可以作为经验值
'''
total_loss = cross_entropy_loss + 7e-5*l2_loss
#添加标量统计结果。
tf.summary.scalar("cross_entropy_loss",cross_entropy_loss)
tf.summary.scalar("l2_loss",l2_loss)
tf.summary.scalar("total_loss",total_loss)
#优化器，选择梯度下降，学习率为0.3
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.3).minimize(total_loss)

pred = tf.nn.softmax(logits)
'''
equal -- 逐个元素判断是否相等，如相等，返回True，否则，False
返回的数据维度与输入一样
argmax -- 返回对应axis维度上最大值的索引；axis--0--表示列上，1--表示行
cast -- 强行转换Tensor类型
'''
correct_pred = tf.equal(tf.argmax(y,1),tf.argmax(logits,1))
accuracy = tf.reduce_mean(tf.cast(correct_pred,tf.float32))

batch_size = 100
training_step = 1100
#tensorflow saver对象
'''
tf.train.Saver() -- Saver对象
saver.save(sess,  '路径 + 模型文件名'）
Tensorflow 会自动生成4个文件：
    （1）model.ckpt.meta -- 保存计算图结构，理解为神经网络的网络结构
    （2）model.ckpt.index和model.ckpt.data -- 保存所有变量取值
    （3）checkpoint -- 保存一个目录下所有模型文件列表
之后可以使用saver.restore()方法，重载模型的参数，继续训练或者用于测试数据 saver.restore(sess,FLAGG.train_dir) 
'''
saver = tf.train.Saver()
merged = tf.summary.merge_all()
with tf.Session() as sess:
    writer = tf.summary.FileWriter("logs/",sess.graph)
    sess.run(tf.global_variables_initializer())

    validate_data = {
        x:mnist.validation.images,
        y:mnist.validation.labels,
        keep_prob:1.0
    }
    test_data = {x:mnist.test.images,y:mnist.test.labels,keep_prob:1.0}

    for i in range(training_step):
        xs,ys = mnist.train.next_batch(batch_size)
        _,loss,rs = sess.run(
            [optimizer,cross_entropy_loss,merged],
            feed_dict={x:xs,y:ys,keep_prob:0.6}
        )
        writer.add_summary(rs,i)

        if i>0 and i%100==0:
            validate_accuracy = sess.run(accuracy,feed_dict=validate_data)
            print("after %d steps,h=the loss is %g,the validation accuracy is %g" % (i,loss,validate_accuracy))
            saver.save(sess,'./model.ckpt',global_step=i)
    print("the training is finished!")
    acc = sess.run(accuracy,feed_dict=test_data)
    print("the test accuracy is :",acc)


with tf.Session() as sess:
    '''
    tf.train.get_checkpoint_state 
    输入参数：路径
    返回的是checkpoint文件proto类型的内容
    返回变量有两个属性，model_checkpoint_path（保存最新的tensorflow模型文件的路径文件名）和all_model_checkpoint_paths（所有未被删除的tensorflow模型文件的路径文件名）
    '''
    ckpy = tf.train.get_checkpoint_state("./")
    if ckpy and ckpy.model_checkpoint_path:
        saver.restore(sess,ckpy.model_checkpoint_path)
        final_pred,acc = sess.run(
            [pred,accuracy],feed_dict={x:mnist.test.images[:16],
                                       y:mnist.test.labels[:16],
                                       keep_prob:1.0}
        )
        '''
        argsort -- 返回的是数组值从小到大的索引值列表
        '''
        orders = np.argsort(final_pred)
        plt.figure(figsize=(8,8))
        print(acc)
        for idx in range(16):
            order = orders[idx,:][-1]
            prob = final_pred[idx,:][order]
            plt.subplot(4,4,idx+1)
            plt.axis('off')
            plt.title('{}:[{}] - [{:.1f}%]'.format(np.argmax(mnist.test.labels[idx]),order,prob*100))
            plt.imshow(mnist.test.images[idx].reshape(28,28))
        plt.show()
    else:
        pass
